package com.flink.hbase.kafka2elasticsearch.timepartition;

import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
import org.elasticsearch.action.admin.indices.alias.IndicesAliasesRequest;
import org.elasticsearch.action.admin.indices.alias.get.GetAliasesRequest;
import org.elasticsearch.action.admin.indices.alias.get.GetAliasesResponse;
import org.elasticsearch.action.admin.indices.create.CreateIndexRequest;
import org.elasticsearch.action.admin.indices.create.CreateIndexResponse;
import org.elasticsearch.action.admin.indices.get.GetIndexRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.indices.GetIndexTemplatesRequest;
import org.elasticsearch.client.indices.GetIndexTemplatesResponse;
import org.elasticsearch.client.indices.IndexTemplatesExistRequest;
import org.elasticsearch.client.indices.PutIndexTemplateRequest;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.xcontent.XContentType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Closeable;
import java.io.IOException;
import java.io.Serializable;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Elasticsearch 索引管理器
 * 负责创建按月分区的索引和管理别名
 */
public class ElasticsearchIndexManager implements Serializable, Closeable {
    private static final Logger LOG = LoggerFactory.getLogger(ElasticsearchIndexManager.class);
    private static final long serialVersionUID = 1L;
    
    private final Map<String, Object> esConfig;
    private final String indexTemplate;
    private final Map<String, String> indexSettings;
    private final Map<String, Object> indexMappings;
    
    private transient RestHighLevelClient client;
    private transient Set<String> createdIndices;
    private transient Set<String> createdAliases;
    
    public ElasticsearchIndexManager(Map<String, Object> esConfig, 
                                   String indexTemplate,
                                   Map<String, String> indexSettings,
                                   Map<String, Object> indexMappings) {
        this.esConfig = esConfig;
        this.indexTemplate = indexTemplate;
        this.indexSettings = indexSettings;
        this.indexMappings = indexMappings;
        this.createdIndices = ConcurrentHashMap.newKeySet();
        this.createdAliases = ConcurrentHashMap.newKeySet();
    }
    
    /**
     * 初始化 ES 客户端和索引模板
     */
    public void initialize() throws IOException {
        initializeElasticsearchClient();
        createIndexTemplate();
    }
    
    /**
     * 确保索引和别名存在
     */
    public void ensureIndexAndAlias(TimePartitionInfo partitionInfo) throws IOException {
        if (client == null) {
            initialize();
        }
        
        String indexName = partitionInfo.getActualIndexName();
        String aliasName = partitionInfo.getAliasName();
        
        // 检查并创建索引
        if (!createdIndices.contains(indexName) && !indexExists(indexName)) {
            createIndex(partitionInfo);
            createdIndices.add(indexName);
            LOG.info("Created index: {}", indexName);
        }
        
        // 检查并创建别名
        String aliasKey = aliasName + ":" + indexName;
        if (!createdAliases.contains(aliasKey) && !aliasExists(aliasName, indexName)) {
            createOrUpdateAlias(partitionInfo);
            createdAliases.add(aliasKey);
            LOG.info("Created/Updated alias: {} -> {}", aliasName, indexName);
        }
    }
    
    /**
     * 检查索引是否存在
     */
    private boolean indexExists(String indexName) throws IOException {
        GetIndexRequest request = new GetIndexRequest(indexName);
        return client.indices().exists(request, RequestOptions.DEFAULT);
    }
    
    /**
     * 检查别名是否存在
     */
    private boolean aliasExists(String aliasName, String indexName) throws IOException {
        GetAliasesRequest request = new GetAliasesRequest(aliasName);
        GetAliasesResponse response = client.indices().getAlias(request, RequestOptions.DEFAULT);
        return response.getAliases().containsKey(indexName);
    }
    
    /**
     * 创建索引
     */
    private void createIndex(TimePartitionInfo partitionInfo) throws IOException {
        String indexName = partitionInfo.getActualIndexName();
        
        CreateIndexRequest request = new CreateIndexRequest(indexName);
        
        // 设置索引配置
        Settings.Builder settingsBuilder = Settings.builder();
        if (indexSettings != null) {
            for (Map.Entry<String, String> entry : indexSettings.entrySet()) {
                settingsBuilder.put(entry.getKey(), entry.getValue());
            }
        }
        
        // 默认设置
        settingsBuilder
            .put("index.number_of_shards", getConfigValue("index.number_of_shards", "3"))
            .put("index.number_of_replicas", getConfigValue("index.number_of_replicas", "1"))
            .put("index.refresh_interval", getConfigValue("index.refresh_interval", "30s"))
            .put("index.max_result_window", getConfigValue("index.max_result_window", "100000"));
        
        request.settings(settingsBuilder);
        
        // 设置映射
        if (indexMappings != null && !indexMappings.isEmpty()) {
            try {
                String mappingJson = new com.fasterxml.jackson.databind.ObjectMapper()
                    .writeValueAsString(indexMappings);
                request.mapping(mappingJson, XContentType.JSON);
            } catch (Exception e) {
                LOG.error("Error setting index mapping", e);
            }
        }
        
        CreateIndexResponse response = client.indices().create(request, RequestOptions.DEFAULT);
        if (!response.isAcknowledged()) {
            throw new IOException("Failed to create index: " + indexName);
        }
    }
    
    /**
     * 创建或更新别名
     */
    private void createOrUpdateAlias(TimePartitionInfo partitionInfo) throws IOException {
        String indexName = partitionInfo.getActualIndexName();
        String aliasName = partitionInfo.getAliasName();
        
        IndicesAliasesRequest request = new IndicesAliasesRequest();
        IndicesAliasesRequest.AliasActions aliasAction = 
            new IndicesAliasesRequest.AliasActions(IndicesAliasesRequest.AliasActions.Type.ADD)
                .index(indexName)
                .alias(aliasName);
        
        request.addAliasAction(aliasAction);
        
        boolean acknowledged = client.indices().updateAliases(request, RequestOptions.DEFAULT).isAcknowledged();
        if (!acknowledged) {
            throw new IOException("Failed to create/update alias: " + aliasName + " -> " + indexName);
        }
    }
    
    /**
     * 创建索引模板
     */
    private void createIndexTemplate() throws IOException {
        if (indexTemplate == null || indexTemplate.trim().isEmpty()) {
            return;
        }
        
        String templateName = indexTemplate + "_template";
        
        // 检查模板是否已存在
        IndexTemplatesExistRequest existsRequest = new IndexTemplatesExistRequest(templateName);
        boolean templateExists = client.indices().existsTemplate(existsRequest, RequestOptions.DEFAULT);
        
        if (templateExists) {
            LOG.info("Index template {} already exists", templateName);
            return;
        }
        
        // 创建索引模板
        PutIndexTemplateRequest request = new PutIndexTemplateRequest(templateName);
        request.patterns(java.util.Arrays.asList(indexTemplate + "_*"));
        
        // 设置模板配置
        Settings.Builder settingsBuilder = Settings.builder();
        if (indexSettings != null) {
            for (Map.Entry<String, String> entry : indexSettings.entrySet()) {
                settingsBuilder.put(entry.getKey(), entry.getValue());
            }
        }
        
        // 默认模板设置
        settingsBuilder
            .put("index.number_of_shards", getConfigValue("template.number_of_shards", "3"))
            .put("index.number_of_replicas", getConfigValue("template.number_of_replicas", "1"))
            .put("index.refresh_interval", getConfigValue("template.refresh_interval", "30s"));
        
        request.settings(settingsBuilder);
        
        // 设置模板映射
        if (indexMappings != null && !indexMappings.isEmpty()) {
            try {
                String mappingJson = new com.fasterxml.jackson.databind.ObjectMapper()
                    .writeValueAsString(indexMappings);
                request.mapping(mappingJson, XContentType.JSON);
            } catch (Exception e) {
                LOG.error("Error setting template mapping", e);
            }
        }
        
        boolean acknowledged = client.indices().putTemplate(request, RequestOptions.DEFAULT).isAcknowledged();
        if (acknowledged) {
            LOG.info("Created index template: {}", templateName);
        } else {
            LOG.warn("Failed to create index template: {}", templateName);
        }
    }
    
    /**
     * 初始化 Elasticsearch 客户端
     */
    private void initializeElasticsearchClient() {
        // 解析主机配置
        String hostsConfig = (String) esConfig.get("hosts");
        String[] hostArray = hostsConfig.split(",");
        HttpHost[] hosts = new HttpHost[hostArray.length];
        
        for (int i = 0; i < hostArray.length; i++) {
            String[] hostPort = hostArray[i].trim().split(":");
            String host = hostPort[0];
            int port = hostPort.length > 1 ? Integer.parseInt(hostPort[1]) : 9200;
            String scheme = (String) esConfig.getOrDefault("scheme", "http");
            hosts[i] = new HttpHost(host, port, scheme);
        }
        
        // 构建客户端
        RestClientBuilder builder = RestClient.builder(hosts);
        
        // 配置认证
        if (esConfig.containsKey("username") && esConfig.containsKey("password")) {
            String username = (String) esConfig.get("username");
            String password = (String) esConfig.get("password");
            
            final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
            credentialsProvider.setCredentials(AuthScope.ANY, 
                new UsernamePasswordCredentials(username, password));
            
            builder.setHttpClientConfigCallback(new RestClientBuilder.HttpClientConfigCallback() {
                @Override
                public HttpAsyncClientBuilder customizeHttpClient(HttpAsyncClientBuilder httpClientBuilder) {
                    return httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
                }
            });
        }
        
        // 配置超时
        builder.setRequestConfigCallback(requestConfigBuilder -> {
            return requestConfigBuilder
                .setConnectTimeout((Integer) esConfig.getOrDefault("connect.timeout", 5000))
                .setSocketTimeout((Integer) esConfig.getOrDefault("socket.timeout", 60000))
                .setConnectionRequestTimeout((Integer) esConfig.getOrDefault("connection.request.timeout", 5000));
        });
        
        client = new RestHighLevelClient(builder);
    }
    
    /**
     * 获取配置值
     */
    private String getConfigValue(String key, String defaultValue) {
        Object value = esConfig.get(key);
        return value != null ? value.toString() : defaultValue;
    }
    
    /**
     * 获取已创建的索引数量
     */
    public int getCreatedIndicesCount() {
        return createdIndices.size();
    }
    
    /**
     * 获取已创建的别名数量
     */
    public int getCreatedAliasesCount() {
        return createdAliases.size();
    }
    
    @Override
    public void close() throws IOException {
        if (client != null) {
            client.close();
        }
    }
    
    /**
     * 构建器
     */
    public static class Builder {
        private Map<String, Object> esConfig;
        private String indexTemplate;
        private Map<String, String> indexSettings;
        private Map<String, Object> indexMappings;
        
        public Builder esConfig(Map<String, Object> esConfig) {
            this.esConfig = esConfig;
            return this;
        }
        
        public Builder indexTemplate(String indexTemplate) {
            this.indexTemplate = indexTemplate;
            return this;
        }
        
        public Builder indexSettings(Map<String, String> indexSettings) {
            this.indexSettings = indexSettings;
            return this;
        }
        
        public Builder indexMappings(Map<String, Object> indexMappings) {
            this.indexMappings = indexMappings;
            return this;
        }
        
        public ElasticsearchIndexManager build() {
            if (esConfig == null) {
                throw new IllegalArgumentException("ES config is required");
            }
            return new ElasticsearchIndexManager(esConfig, indexTemplate, indexSettings, indexMappings);
        }
    }
    
    /**
     * 创建构建器
     */
    public static Builder builder() {
        return new Builder();
    }
} 